import numpy as np
import os
import json
import faiss
import struct


def ivecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('i' * len(x), *x))

    f.close()


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def fvecs_read(fname):
    data, d = ivecs_read(fname)
    return data.view('float32').astype(np.float32), d


def ip_gnd(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatIP(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    return gnd_idx, gnd_distance


def get_info(dataset_name, read_basic_dir, topk, new_basic_dir):
    base, dim = fvecs_read("%s/%s/%s_base.fvecs" % (read_basic_dir, dataset_name, dataset_name))
    query, dim = fvecs_read("%s/%s/%s_query.fvecs" % (read_basic_dir, dataset_name, dataset_name))
    np.savetxt('%s/%s/%s_query.txt' % (new_basic_dir, dataset_name, dataset_name), query)

    base_len = len(base)
    query_len = len(query)
    dimension = dim

    dataset_info = {
        "vecsize": str(base_len),
        "vecdim": str(dimension),
        "qsize": str(query_len),
        "dataset": dataset_name
    }
    index_cmd = '../../build/build_graph.sh %s_base.txt %d %d' % (
        dataset_name, base_len, dimension)
    query_cmd = "../../build/test_query.sh %s_query.txt %d %d %d %d gnd.ivecs" % (
        dataset_name, base_len, dimension, 100, topk)
    print("get info dataset %s" % dataset_name)
    return dataset_info, index_cmd, query_cmd


'''
sh地址 base文件名 base的长度 数据集维度
sh地址 query文件名 base的长度 数据集维度 n_candidate(任意值) topk gnd文件名

../../build_graph.sh audio_base.txt 53387 192

../../test_query.sh audio_query.txt 53387 192 20 10 gnd.ivecs
'''

if __name__ == '__main__':
    # basic_dir = '/home/zhengbian/Dataset/inner_product'
    basic_dir = '/home/zhengbian/plus-ip-nsw/data'
    new_basic_dir = '/home/zhengbian/mobius/data'
    dataset_name_l = ['audio', 'imagenet', 'movielens', 'music100', 'netflix', 'normal-64', 'text-to-image', 'tiny5m',
                      'word2vec', 'yahoomusic']
    topk = 10
    ds_info_l = []

    dataset_info_m = {}
    index_cmd_m = {}
    query_cmd_m = {}

    for ds_name in dataset_name_l:
        dataset_info, tmp_index_cmd, tmp_query_cmd_l = get_info(ds_name, basic_dir, topk, new_basic_dir)
        dataset_info_m[ds_name] = dataset_info
        index_cmd_m[ds_name] = tmp_index_cmd
        query_cmd_m[ds_name] = tmp_query_cmd_l

    with open("dataset_info.json", 'w') as f:
        json.dump(dataset_info_m, f)
    with open("index_cmd.json", 'w') as f:
        json.dump(index_cmd_m, f)
    with open("query_cmd.json", 'w') as f:
        json.dump(query_cmd_m, f)
